from PIL import Image


def make_image_even(image, n):
    """
    取得一个PIL图像并且更改所有值为偶数（使最低有效位为0）
    """
    pixels = list(image.getdata())  # 得到一个这样的列表：[(r,g,b,t),(r,g,b,t)...]
    even_pixels = [(r >> n << n, g >> n << n, b >> n << n, t >> n << n) for [r, g, b, t] in pixels]  # 列表推导式，更改所有值为偶数
    even_image = Image.new(image.mode, image.size)  # 创建一个相同大小的图片副本
    even_image.putdata(even_pixels)  # 把上面的像素放入到图片副本
    return even_image


def const_len_bin(int):
    """
    内置函数bin()的替代，返回固定长度的二进制字符串
    """
    binary = "0" * (8 - (len(bin(int)) - 2)) + bin(int).replace('0b', '')  # 去掉bin()返回的二进制字符串中的'0b'，并在左边补足'0'直到字符串长度为8
    return binary


def encode_data_in_image(image, data):
    """
    采用LSB将数据隐写进图片中
    """
    even_image = make_image_even(image)  # 获得最低有效位为0的图片副本
    binary = ''.join(map(const_len_bin, bytearray(data, 'utf-8')))  # 将需要被隐藏的字符串转换成二进制字符串
    if len(binary) > len(image.getdata()) * 4:  # 如果不可能编码全部数据，抛出异常。
        raise Exception("Error: Can't encode more than " + len(even_image.getdata()) * 4 + " bits in this image.")
    encoded_pixels = [
        (
            r + int(binary[index * 4 + 0]),
            g + int(binary[index * 4 + 1]),
            b + int(binary[index * 4 + 2]),
            t + int(binary[index * 4 + 3])
        ) if index * 4 < len(binary) else (r, g, b, t) for index, (r, g, b, t) in enumerate(list(even_image.getdata()))
    ]  # 将binary中的二进制字符串信息编码进像素里
    encoded_image = Image.new(even_image.mode, even_image.size)  # 创建新图片以存放编码后的像素
    encoded_image.putdata(encoded_pixels)  # 添加编码后的数据
    return encoded_image


def binary_to_string(binary):
    index = 0
    string = []
    rec = lambda x, i: x[2:8] + (rec(x[8:], i-1) if i > 1 else '') if x else ''
    fun = lambda x, i: x[i+1:8] + rec(x[8:], i-1)
    while index + 1 < len(binary):
        chartype = binary[index:].index('0')  # 存放字符所占字节数，一个字节的字符会存为0
        length = chartype * 8 if chartype else 8
        string.append(chr(int(fun(binary[index:index+length],chartype), 2)))
        index += length
    return ''.join(string)


def decode_image(image):
    """
    从图片中解码数据
    """
    pixels = list(image.getdata())  # 获得编码后的像素列表
    binary = ''.join([str(int(r >> 1 << 1 != r)) + str(int(g >> 1 << 1 != g)) + str(int(b >> 1 << 1 != b)) + str(int(t >> 1 << 1 != t)) for (r, g, b, t) in pixels])  # 提取图片中所有最低有效位中的数据
    # 找到数据截止处的标志，返回索引
    location_double_null = binary.find('0000000000000000')
    # 不足位补零
    end_index = location_double_null + (8 - (location_double_null % 8)) if location_double_null % 8 != 0 else location_double_null
    data = binary_to_string(binary[0: end_index])
    return data


if __name__ == '__main__':
    encode_data_in_image(Image.open('images/before.png').convert("RGBA"), '你好世界').save('images/after.png')
    print(decode_image(Image.open('images/after.png')))
